import torch
from torch.utils.data import DataLoader

from DataSetLoader import *
from model import *

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# 加载测试数据集
test_dataset = ECGDataSet("./DataSetProcess/MyFFTValidation2017/")
test_loader = DataLoader(test_dataset, batch_size=1, shuffle=True)

# 创建模型实例并加载训练好的权重
rnn_model = LSTMClassifier(5120, 360, 4, 4)
rnn_model.load_state_dict(torch.load('bestNetv2.pth'))
rnn_model.to(device)

# 设置模型为评估模式
rnn_model.eval()
LableDecode = ["N", "A", "O", "~"]
# 进行预测
with torch.no_grad():
    correct = 0
    total = 0
    for data in test_loader:
        ecg, labels = data
        ecg = ecg.to(device)
        labels = labels.to(device)
        outputs = rnn_model(ecg)
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()

    accuracy = 100 * correct / total
    print("Test Accuracy: {:.2f}%".format(accuracy))